package com.xh.shuati.ai.rag;

import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.document.splitter.DocumentByParagraphSplitter;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.store.embedding.EmbeddingStore;
import dev.langchain4j.store.embedding.EmbeddingStoreIngestor;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import javax.annotation.Resource;
import java.util.List;

/**
 *  rag 内容增强器配置类
 */
@Configuration
public class RagConfig {

    @Resource
    private EmbeddingModel myQwenEmbeddingModel;

    @Resource
    EmbeddingStore<TextSegment> myEmbeddingStore;

    @Bean
    public ContentRetriever myContentRetriever() {

        /**
        //文档加载
        List<Document> documents = FileSystemDocumentLoader.loadDocuments("src/main/resources/docs");
        //自定义分词器
        DocumentByParagraphSplitter paragraphSplitter = new DocumentByParagraphSplitter(500, 100);
        //加载文档
        // 3. 自定义文档加载器
        EmbeddingStoreIngestor ingestor = EmbeddingStoreIngestor.builder()
                .documentSplitter(paragraphSplitter)
                // 为了提高搜索质量，为每个 TextSegment 添加文档名称
                .textSegmentTransformer(textSegment -> TextSegment.from(
                        textSegment.metadata().getString("file_name") + "\n" + textSegment.text(),
                        textSegment.metadata()
                ))
                // 使用指定的向量模型
                .embeddingModel(myQwenEmbeddingModel)
                .embeddingStore(myEmbeddingStore)
                .build();

        // 加载文档
        ingestor.ingest(documents);
         **/

        // 4. 自定义内容查询器
        ContentRetriever contentRetriever = EmbeddingStoreContentRetriever.builder()
                .embeddingStore(myEmbeddingStore)
                .embeddingModel(myQwenEmbeddingModel)
                .maxResults(5) // 最多 5 个检索结果
                .minScore(0.5) // 过滤掉分数小于 0.5 的结果
                .build();
        return contentRetriever;
    }
}
